# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: donggu

#
# ❯ ../run_exp.py python/coin_deploy/lm_agg2/config_coin1_style/donggu/lm_agg_okex_futures_allexpiry_ltc.json LTC-USD.QUARTER 20200322-20200326 24H

import subprocess
import json
import optuna
import sys
import copy
import tempfile
import re
import math
import socket
import neptune
import neptunecontrib.monitoring.optuna as opt_utils
from coin2.strategy.pta.performance import donggu2


class Search(object):
  def __init__(self,
               exp_group,
               exp_name,
               callback_set_exp_property,
               callback_set_params_and_cmds,
               n_trials=200,
               n_jobs=1,
               eod_clear=True):
    """[summary]

    Args:
        exp_group: e.g. 'lm-agg'
        exp_name: e.g. 'lm-agg-BTC-USD.PERPETUAL-20200501-20200503-24H'
        callback_set_exp_property: a function to set neptune exp properties
        callback_set_params_and_cmds: a function to set params and returns cmds
        eod_clear: pnl calculation assumes eod-clear (position to 0)
    """
    self.exp_group = exp_group
    self.exp_name = exp_name
    self.set_exp_property = callback_set_exp_property
    self.set_params_and_cmds = callback_set_params_and_cmds
    self.n_trials = n_trials
    self.n_jobs = n_jobs
    self.eod_clear = eod_clear

    neptune.init(f'donggu-prestolabs/{exp_group}')
    self.neptune_exp = neptune.create_experiment(name=exp_name)
    self.set_exp_property(self.neptune_exp)

  def run_trial(self, trial):
    cmds = self.setup_params_and_cmds(trial)
    assert isinstance(cmds, list)
    print(f"Trial #{trial.number} cmd: {' '.join(cmds)}")
    out = subprocess.check_output(cmds, shell=False, stderr=subprocess.DEVNULL)

    # capture daily pnls
    try:
      reports = list(map(json.loads, out.splitlines()))  # multi dates => 1 json report per line
    except Exception:
      print('error parsing daily pnl report: ', out)
      raise

    pnls = []
    assert len(reports) >= 1
    try:
      last_price = 0
      ts = 0
      for report in reports:
        for (k, v) in report.items():  # only single product
          t = v['last_ts']
          if t > ts and v['last_price_in_settle_currency'] > 0:
            ts = t
            last_price = v['last_price_in_settle_currency']
      for report in reports:
        for (k, v) in report.items():
          if self.eod_clear or v['last_price_in_settle_currency'] == 0:
            unrealized_scale = 0
          else:
            unrealized_scale = (last_price / v['last_price_in_settle_currency'] - 1.0)

          # trading once per 5 mins => 288 per day
          # once per 1 min => 1440
          fill_incentive = min(v['num_fills'], 288) / 288.0  # 0~1
          if v['pnl_net'] < 0:
            # + pnl? fill_incentive = 0~1.0
            # - pnl? fill_incentive = 2.0~1.0
            fill_incentive = 2 - fill_incentive

          pnls.append((v['pnl_net'] + v['pnl_unrealized'] * unrealized_scale) * fill_incentive)
    except Exception:
      print(f'Trial #{trial.number}: error reading pnl report: ', reports)
      raise

    score = donggu2(pnls)
    print(f'Trial #{trial.number}: score {score} / pnls {pnls}')

    # upload pnl plot
    self.neptune_exp.log_image('pnl_plot',
                               x=trial.number,
                               y=f'out/{self.exp_name}/trial{trial.number}_total.png',
                               name=f'Trial {trial.number}')
    self.neptune_exp.log_metric('pnl', x=trial.number, y=score)

    return score

  def run(self, epoch):
    study = optuna.create_study(study_name=self.exp_name,
                                direction='maximize',
                                storage=f'postgresql://postgres:EU2kb6Xos9pS@10.21.105.13:5432',
                                load_if_exists=True)
    print(f'search exp name = {self.exp_name}')
    print(f'hostname = {socket.gethostname()}')
    print('==== start ====')
    study.optimize(self.run_trial,
                   n_trials=self.n_trials,
                   n_jobs=self.n_jobs,
                   callbacks=[opt_utils.NeptuneMonitor()])
    opt_utils.log_study(study)

    print('==== search result ====')
    print('best  score:', study.best_score)
    print('best params:', study.best_params)
